/**
 * OWASP Enterprise Security API (ESAPI)
 *
 * This file is part of the Open Web Application Security Project (OWASP)
 * Enterprise Security API (ESAPI) project. For details, please see
 * <a href="http://www.owasp.org/index.php/ESAPI">http://www.owasp.org/index.php/ESAPI</a>.
 *
 * Copyright (c) 2009 - The OWASP Foundation
 *
 * The ESAPI is published by OWASP under the BSD license. You should read and accept the
 * LICENSE before you use, modify, and/or redistribute this software.
 *
 * @author Arshan Dabirsiaghi <a href="http://www.aspectsecurity.com">Aspect Security</a>
 * @created 2009
 */
package org.owasp.esapi.waf.internal;

import java.io.File;
import java.io.FileNotFoundException;
import java.io.IOException;
import java.io.RandomAccessFile;

import javax.servlet.ServletOutputStream;
import javax.servlet.WriteListener;

/**
 * This class was inspired by ModSecurity for Java by Ivan Ristic. We hook
 * the response stream and queue up all outbound data so that we can apply
 * egress rules. For efficiency, we decide off the bat if we need to buffer
 * responses to accomplish any of the rules in the policy file.
 *
 * If not, we just forward everything through, otherwise we write data to our
 * byte stream that we will eventually forward en totale to the user agent.
 *
 * @author Arshan Dabirsiaghi
 */

public class InterceptingServletOutputStream extends ServletOutputStream {

    private static final int FLUSH_BLOCK_SIZE = 1024;
    private ServletOutputStream os;
    private boolean buffering;
    private boolean committed;
    private boolean closed;

    private RandomAccessFile out;

    public InterceptingServletOutputStream(ServletOutputStream os, boolean buffered) throws FileNotFoundException, IOException {
        super();
        this.os = os;
        this.buffering = buffered;
        this.committed = false;
        this.closed = false;

        /*
         * Creating a RandomAccessFile to keep track of output generated. I made
         * the prefix and suffix small for less processing. The "oew" is intended
         * to stand for "OWASP ESAPI WAF" and the "hop" for HTTP output.
         */
        File tempFile= File.createTempFile("oew", ".hop");
        this.out = new RandomAccessFile (tempFile, "rw" );
        tempFile.deleteOnExit();

    }

    public void reset() throws IOException {
        out.setLength(0L);
    }

    public byte[] getResponseBytes() throws IOException {

        byte[] buffer = new byte[(int) out.length()];
        out.seek(0);
        out.read(buffer, 0, (int)out.length());
        out.seek(out.length());
        return buffer;

    }

    public void setResponseBytes(byte[] responseBytes) throws IOException {

        if ( ! buffering && out.length() > 0 ) {
            throw new IOException("Already committed response because not currently buffering");
        }

        out.setLength(0L);
        out.write(responseBytes);
    }

    public void write(int i) throws IOException {
        if (!buffering) {
            os.write(i);
        }
        out.write(i);
    }

    public void write(byte[] b) throws IOException {
        if (!buffering) {
            os.write(b, 0, b.length);
        }
        out.write(b, 0, b.length);
    }

    public void write(byte[] b, int off, int len) throws IOException {
        if (!buffering) {
            os.write(b, off, len);
        }
        out.write(b, off, len);
    }

    public void flush() throws IOException {

        if (buffering) {

            synchronized(out) {

                out.seek(0);

                byte[] buff = new byte[FLUSH_BLOCK_SIZE];

                for(int i=0;i<out.length();) {

                    long currentPos = out.getFilePointer();
                    long totalSize = out.length();
                    int amountToWrite = FLUSH_BLOCK_SIZE;

                    if ( (totalSize - currentPos) < FLUSH_BLOCK_SIZE ) {
                        amountToWrite = (int) (totalSize - currentPos);
                    }

                    out.read(buff, 0, amountToWrite);

                    os.write(buff,0,amountToWrite);

                    i+=amountToWrite;

                }

                out.setLength(0);

            }
        }

    }

    public void commit() throws IOException {

        if (!buffering) { // || committed || closed
            return;
        } else {
            flush();
        }
        committed = true;
    }

    public void close() throws IOException {

        if (!buffering)  {
            os.close();
        }
        closed = true;

    }

    @Override
    public boolean isReady() {
        return os.isReady();
    }

    @Override
    public void setWriteListener(WriteListener writeListener) {
        os.setWriteListener(writeListener);
    }

}
